import torch  # 导入PyTorch库

# 定义模型参数
input_size = 4  # 输入特征维度（字符总数）
hidden_size = 4  # 隐藏层维度
num_layers = 1  # RNN层数
batch_size = 1  # 批次大小
seq_len = 5  # 序列长度（每个样本包含的字符数）

# 准备数据
idx2char = ['e', 'h', 'l', 'o']  # 字符到索引的映射：0->e, 1->h, 2->l, 3->o
x_data = [1, 0, 2, 2, 3]  # 输入序列的索引：对应"h e l l o"
y_data = [3, 1, 2, 3, 2]  # 目标序列的索引：对应"o h l o l"

# 创建one-hot编码查找表
# 每个字符用一个4维向量表示，对应位置为1，其余为0
one_hot_lookup = [[1, 0, 0, 0],  # 对应索引0 ('e')
                  [0, 1, 0, 0],  # 对应索引1 ('h')
                  [0, 0, 1, 0],  # 对应索引2 ('l')
                  [0, 0, 0, 1]]  # 对应索引3 ('o')

# 将输入序列转换为one-hot编码
x_one_hot = [one_hot_lookup[x] for x in x_data]  # 结果是5个4维向量的列表
print('x_one_hot:', x_one_hot)

# 构造输入张量和标签张量
# 将输入转换为形状为(seq_len, batch_size, input_size)的张量
inputs = torch.Tensor(x_one_hot).view(seq_len, batch_size, input_size)
print(inputs)
# 标签保持为长整型张量，形状为(seq_len,)
labels = torch.LongTensor(y_data)
print(labels)


# 设计RNN模型
class Model(torch.nn.Module):
    def __init__(self, input_size, hidden_size, batch_size, num_layers=1):
        super(Model, self).__init__()
        self.num_layers = num_layers  # RNN层数
        self.batch_size = batch_size  # 批次大小
        self.input_size = input_size  # 输入特征维度
        self.hidden_size = hidden_size  # 隐藏层维度

        # 定义RNN层
        self.rnn = torch.nn.RNN(
            input_size=self.input_size,  # 输入特征维度
            hidden_size=self.hidden_size,  # 隐藏层维度
            num_layers=self.num_layers  # 层数
        )

    def forward(self, input):
        # 初始化隐藏状态为全零
        hidden = torch.zeros(self.num_layers, self.batch_size, self.hidden_size)

        # RNN前向计算：out包含所有时间步的输出
        out, _ = self.rnn(input, hidden)

        # 重塑输出形状为(seq_len*batch_size, hidden_size)，以便计算交叉熵损失
        return out.view(-1, self.hidden_size)


# 创建模型实例
net = Model(input_size, hidden_size, batch_size, num_layers)

# 定义损失函数和优化器
criterion = torch.nn.CrossEntropyLoss()  # 交叉熵损失，适用于分类问题
optimizer = torch.optim.Adam(net.parameters(), lr=0.05)  # Adam优化器，学习率0.05

# 训练循环
for epoch in range(20):  # 训练20轮
    optimizer.zero_grad()  # 清零梯度

    # 前向传播：获取模型输出
    outputs = net(inputs)  # outputs形状: (seq_len*batch_size, hidden_size)

    # 计算损失
    loss = criterion(outputs, labels)

    # 反向传播与参数更新
    loss.backward()  # 计算梯度
    optimizer.step()  # 更新参数

    # 计算预测结果（取概率最大的索引）
    _, idx = outputs.max(dim=1)  # idx是预测的字符索引
    idx = idx.data.numpy()  # 转换为numpy数组

    # 打印预测结果和损失
    print('Predicted: ', ''.join([idx2char[x] for x in idx]), end='')  # 将索引转换为字符
    print(', Epoch [%d/20] loss=%.3f' % (epoch + 1, loss.item()))  # 打印轮次和损失